# -*- coding: utf-8 -*-
###################################################
#    This file is part of blockIMH.
#
#    blockIMH is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    blockIMH is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with blockIMH.  If not, see <http://www.gnu.org/licenses/>.
###################################################
#! /usr/bin/env python

from __future__ import division
import os
import sys
import scipy.weave as weave
from scipy.stats import norm, cauchy, uniform
from scipy import special
from numpy import random, exp, sqrt, power, log, newaxis, zeros, \
        zeros_like, float32, int32, repeat, minimum, ones, mean, average, sum, \
        min, prod, nan, load, vstack
try:
    import pycuda
    import pycuda.autoinit
    import pycuda.gpuarray as gpuarray
    import pycuda.driver as cuda
    import pycuda.cumath as cumath
    from pycuda.compiler import SourceModule
    from mycudafunctions import *
except:
    CUDA_ENABLED = False
import permutations
import raoblackwell

## Target = Normal(0, 1)
## Proposal = Cauchy(0, 1)
dtarget = norm.pdf
dproposal = cauchy.pdf
rproposal = cauchy.rvs

def omega(x):
    """
    Computes omega ratios as
    defined in the article.
    """
    return dtarget(x) / dproposal(x)
def fastomega(x):
    """
    Computes omega ratios as
    defined in the article.
    (Much faster)
    """
    xsquared = x**2
    return 1.253314 * exp(-xsquared/2) * (1 + xsquared)

def simpleIMH(T, x_starting):
    """
    This creates a standard IMH chain. Not actually
    used anywhere else in the code.
    """
    wcurrent = fastomega(x_starting) 
    chain = zeros(T + 1)
    chain[0] = x_starting
    currentidx = 0
    acceptations = zeros(T)
    for t in xrange(T):
        if t % 1000 == 0:
            print "iteration %i" % t
        y = rproposal(size = 1)
        omegaproposal = fastomega(y)
        acceptratio = (omegaproposal / wcurrent)
        u = random.uniform(size = 1)
        acceptations[t] = (u < acceptratio)
        if acceptations[t]:
            chain[t+1] = y
            wcurrent = omegaproposal
        else:
            chain[t+1] = chain[t]
    print "accept rate: %.3f" % (sum(acceptations) / T)
    return chain

def IMHsteps(p, omegaproposalsblock, omegacurrent, uniformsblock, weightsRB1, nboccurs, chain, proposalsblock):
    """
    scipy.weave implementation of a (p * p) block
    """
    code = \
    """
    float acceptratio;
    int currentidx;
    int acceptation;
    float wcurrent;
    for(int k = 0; k < p; k++)
    {
        currentidx = 0;
        wcurrent = omegacurrent;
        for(int t = 0; t < p; t++)
        {
            acceptratio = (double) omegaproposalsblock(t, k) / (double) wcurrent;
            if (acceptratio > 1){ acceptratio = 1;}
            acceptation = (uniformsblock(t, k) < acceptratio);
            weightsRB1(currentidx, k) += 1 - acceptratio;
            weightsRB1(t + 1, k) += acceptratio;
            if (acceptation){
                nboccurs(t + 1, k) += 1;
                currentidx = t + 1 ;
                chain(t + 1, k) = proposalsblock(t, k);
                wcurrent = omegaproposalsblock(t, k);
            }
            else{
                nboccurs(currentidx, k) += 1;
                chain(t+1, k) = chain(t, k);
            }

        }
    }
    """
    weave.inline(code,["p", "omegaproposalsblock", "omegacurrent", "uniformsblock", "weightsRB1", \
            "nboccurs", "chain", "proposalsblock"], type_converters=weave.converters.blitz)

def parallelchains(h, p, proposals, permutations, uniformsblock, omegaproposals, x_starting, omegacurrent, RB1 = False, RB2 = False):
    """
    Runs the block IMH algorithm for one block,
    given proposals, omega rations,
    a permutation scheme, a starting value,
    a function of interest h,...
    """
    # allocate memory for the resulting values 
    chain = zeros((p+1, p))
    chain[0, :] = x_starting
    # permutate what needs to be permutated
    proposalsblock = proposals[permutations]
    omegaproposalsblock = omegaproposals[permutations]
    diffvalues = zeros((p+1, p))
    diffvalues[0, :] = repeat(x_starting, p)
    diffvalues[1:(p+1), :] = proposalsblock
    # weights associated with the first RB method
    weightsRB1 = zeros((p+1, p))
    # occurrences of a proposal withing a block
    nboccurs = zeros((p+1, p))
    # runs a block
    IMHsteps(p, omegaproposalsblock, omegacurrent, uniformsblock, weightsRB1, nboccurs, chain, proposalsblock)
    # this is hat tau 2
    blockestim = average(h(diffvalues), weights = nboccurs)
    # this is hat tau 3
    estimRB1 = average(h(diffvalues), weights = weightsRB1)
    if RB2:
        weightsRB2 = zeros((p+1, p))
        for k in range(p):
            raoblackwell.raoblackwell(k, p, omegaproposalsblock, omegacurrent, weightsRB2)
        # this is hat tau 4
        estimRB2 = average(h(diffvalues), weights = weightsRB2)
    else:
        estimRB2 = 0
    return({"blockestim": blockestim, "estimRB1": estimRB1, "acceptedvalues": chain[1:(p+1), :], "estimRB2" : estimRB2})

def blockIMH(x_starting, T, p, h):
    """
    Runs the whole thing
    """
    x = zeros(T+1)
    x[0] = x_starting
    omegacurrent = fastomega(x_starting)
    nbblocks = T // p
    acceptedvalues = zeros((p, p))
    uniformsblock = zeros((p,p))
    # various permutations
    noPermutations = permutations.createNoPermutation(p)
    circPermutations = permutations.createCircPermutation(p)
    tempPermutations = permutations.createTempPermutation(p)
    # lots of pre allocation
    blockestimRandPermut = zeros(nbblocks)
    estimRB1RandPermut = zeros(nbblocks)
    estimRB2RandPermut = zeros(nbblocks)
    blockestimNoPermut = zeros(nbblocks)
    estimRB1NoPermut = zeros(nbblocks)
    blockestimCircPermut = zeros(nbblocks)
    estimRB1CircPermut = zeros(nbblocks)
    blockestimHalfRandPermut = zeros(nbblocks)
    estimRB1HalfRandPermut = zeros(nbblocks)
    blockestimStratRandPermut = zeros(nbblocks)
    estimRB1StratRandPermut = zeros(nbblocks)
    blockestimTempPermut = zeros(nbblocks)
    # Generate proposals; this could *really* benefit
    # from parallel processors
    allproposals = rproposal(size = p * nbblocks)
    # Compute omega ratios; this could *really* benefit
    # from parallel processors
    allomegaproposals = fastomega(allproposals)
    estimIS = average(h(allproposals), weights = allomegaproposals)
    for block in range(nbblocks):
        # random permutations are drawn for each block
        randPermutations = permutations.createRandomPermutation(p)
        halfRandPermutations = permutations.createHalfRandom(p)
        stratPermutations = permutations.createStratRandom(p)
        #print "block%i" % block
        y = allproposals[(block * p):((block + 1) * p)]
        omegaproposals = allomegaproposals[(block * p):((block + 1) * p)]
        uniformsblock[...] = random.uniform(size = p**2).reshape(p,p)
## Random order
        resultsRandPermut = parallelchains(h, p, y, randPermutations, uniformsblock, \
                omegaproposals, x_starting, omegacurrent, RB1 = True, RB2 = True)
        blockestimRandPermut[block] = resultsRandPermut["blockestim"]
        estimRB1RandPermut[block] = resultsRandPermut["estimRB1"]
        estimRB2RandPermut[block] = resultsRandPermut["estimRB2"]
## Same order
        resultsNoPermut = parallelchains(h, p, y, noPermutations, uniformsblock, omegaproposals, x_starting, omegacurrent)
        blockestimNoPermut[block] = resultsNoPermut["blockestim"]
        estimRB1NoPermut[block] = resultsNoPermut["estimRB1"]
## Circular permutation
        resultsCircPermut = parallelchains(h, p, y, circPermutations, uniformsblock, omegaproposals, x_starting, omegacurrent)
        blockestimCircPermut[block] = resultsCircPermut["blockestim"]
        estimRB1CircPermut[block] = resultsCircPermut["estimRB1"]
## Half Random permutation
        resultsHalfRandPermut = parallelchains(h, p, y, halfRandPermutations, uniformsblock, omegaproposals, x_starting, omegacurrent)
        blockestimHalfRandPermut[block] = resultsHalfRandPermut["blockestim"]
        estimRB1HalfRandPermut[block] = resultsHalfRandPermut["estimRB1"]
## Stratified random permutation
        resultsStratRandPermut = parallelchains(h, p, y, stratPermutations, uniformsblock, omegaproposals, x_starting, omegacurrent)
        blockestimStratRandPermut[block] = resultsStratRandPermut["blockestim"]
        estimRB1StratRandPermut[block] = resultsStratRandPermut["estimRB1"]
## Temp permutation
        resultsTempPermut = parallelchains(h, p, y, tempPermutations, uniformsblock, omegaproposals, x_starting, omegacurrent)
        blockestimTempPermut[block] = resultsTempPermut["blockestim"]
        #estimRB1TempPermut[block] = resultsTempPermut["estimRB1"]
## Prepare the next block by picking an element randomly; from the random permutation scheme.
        selected = int32(random.uniform(size = 1) * p)
        x[(p * block + 1):(p * (block+1) +1)] = resultsRandPermut["acceptedvalues"][:, selected][:, 0]
        x_starting = x[p * (block + 1)]
        omegacurrent = fastomega(x_starting)
    chain = x[1:(T+1)]
    return {"chain": chain, "basicestim": mean(h(chain)), \
            "blockestimRandPermut": mean(blockestimRandPermut), \
            "estimRB1RandPermut": mean(estimRB1RandPermut), \
            "estimRB2RandPermut": mean(estimRB2RandPermut), \
            "blockestimNoPermut": mean(blockestimNoPermut), \
            "estimRB1NoPermut": mean(estimRB1NoPermut), \
            "blockestimCircPermut": mean(blockestimCircPermut), \
            "estimRB1CircPermut": mean(estimRB1CircPermut), \
            "blockestimHalfRandPermut": mean(blockestimHalfRandPermut), \
            "estimRB1HalfRandPermut": mean(estimRB1HalfRandPermut), \
            "blockestimStratRandPermut": mean(blockestimStratRandPermut), \
            "estimRB1StratRandPermut": mean(estimRB1StratRandPermut), \
            "blockestimTempPermut": mean(blockestimTempPermut), \
            "estimIS": estimIS}









